import torch.nn
from torch import nn
import torch.nn.functional as F


class Cos_Criterion(torch.nn.Module):
    """
        余弦损失函数
    """

    def __init__(self):
        super(Cos_Criterion, self).__init__()
        self.cosineSimilarity = nn.CosineSimilarity(dim=1, eps=1e-6)

    def forward(self, x, target):
        # 计算余弦相似度
        similarity = self.cosineSimilarity(x, target)
        # 计算损失
        loss = (1 - similarity).mean()
        return loss


# 计算SVM损失
# https://zhuanlan.zhihu.com/p/343037852
def hinge_loss(outputs, labels):
    """
    折页损失计算
    :param outputs: 大小为(N, num_classes)
    :param labels: 大小为(N)
    :return: 损失值
    """
    num_labels = len(labels)
    #
    corrects = outputs[range(num_labels), labels.squeeze()].unsqueeze(0).T

    # 最大间隔
    margin = 1.0
    margins = outputs - corrects + margin
    loss = torch.sum(torch.max(margins, 1)[0]) / len(labels)

    # # 正则化强度
    # reg = 1e-3
    # loss += reg * torch.sum(weight ** 2)

    return loss


# 计算SVM损失
def svm_loss(scores, label):
    """ https://zhuanlan.zhihu.com/p/343037852
    """
    label2 = label.clone()
    label2[label == 0] = torch.tensor([-1], dtype=torch.long)
    # SVM损失函数，label取值集合为{-1, 1}
    loss = 1 - label2 * scores
    max = loss[torch.gt(loss,0)]
    return torch.sum(max)
